import scipy.ndimage
import torch
import torch.utils
from torch.utils.tensorboard import SummaryWriter
import numpy as np
from pathlib import Path
from datetime import datetime
from sklearn.metrics import balanced_accuracy_score
from sklearn.utils.class_weight import compute_class_weight
from time import time
from tqdm import tqdm
import scipy

from models.transformer import ECGTransformer
from models.cnn import CNN2D
from training.utils import load_chunk_of_data, create_log_file, log_results, NoamOpt
from data.hdf5_dataset import HDF5Dataset

from typing import Tuple, Dict, Any
from torch import Tensor

import warnings
import argparse
import json
warnings.simplefilter("ignore")

torch.set_num_interop_threads(12)
torch.set_num_threads(12)


def tqdm(arg):
    return arg


config = {
    # Transformer parameters
    'd_embedding': 256,
    'n_layers': 2,
    'dropout': .3,
    'with_time_encoding': True,
    'late_fusion': True,

    # Criterion parameters
    'criterion': torch.nn.CrossEntropyLoss,
    'weighted_loss': False,

    # Optimizer parameters
    'optimizer': torch.optim.Adam,
    'warmup': True,
    'lr': 1E-4,
    'regularization': 1E-2,
    'batch_size': 128,

    # General parameters
    'fold_id': 0,
    # 'data_path': Path('/home/timodw/IDLab/Digihealth-Asia/cardiovascular_monitoring/ugent/heartbeat_classification/processed_data/training_snapshots') / 'incart_stratified_standard_25hz_5sr',
    # 'data_path': Path('/home/timodw/IDLab/Digihealth-Asia/cardiovascular_monitoring/ugent/heartbeat_classification/processed_data/training_snapshots') / 'incart_stratified_standard_train_100-125hz_val_500hz',
    'data_path': Path('/home/timodw/IDLab/Digihealth-Asia/cardiovascular_monitoring/ugent/heartbeat_classification/processed_data/training_snapshots') / 'incart_stratified_standard_original',
    'checkpoint_root': Path('/home/timodw/IDLab/Digihealth-Asia/cardiovascular_monitoring/ugent/heartbeat_classification/logs'),

    'checkpoint_name': None,
    'allowed_labels': (0, 1, 2),
    'input_type': 'ecg+cwt+rr', # cwt+rr, cwt, ecg+rr, ecg, cwt+ecg+rr, cwt+ecg
    
    'n_input': 1000,
    'device': 'cuda',
    'chunk_size': 5_000,
    'n_epochs_per_chunk': 10,
    'balanced_chunk': True,
    'n_chunks': 20,
    'positional_encoding_scheme': 'timestamp_concat_after',
    'model': 'transformer',
    'n_iter': 5,
    'patience': 200
}


# class MyTensorboard(TensorBoard):
#     def on_epoch_end(self, net, **kwargs):
#         epoch = net.history[-1, 'epoch']
#         train_loss = net.history[-1, 'train_loss']
#         valid_acc = net.history[-1, 'valid_acc']
#         valid_acc_bal = net.history[-1, 'valid_acc_bal']
#         valid_loss = net.history[-1, 'valid_loss']
#         self.writer.add_scalar('accuracy/valid_acc', valid_acc, epoch)
#         self.writer.add_scalar('accuracy/valid_acc_bal', valid_acc_bal, epoch)

#         self.writer.add_scalar('loss/training', train_loss, epoch)
#         self.writer.add_scalar('loss/validation', valid_loss, epoch)


def process_chunk(model: ECGTransformer, criterion: torch.nn.modules.loss._Loss,
                  optimizer: torch.optim.Optimizer,
                  training_dataloader: torch.utils.data.DataLoader,
                  validation_dataloader: torch.utils.data.DataLoader,
                  config: Dict[str, Any], chunk_id: int, current_patience: int, best_loss=float('inf')):
    for it in range(config['n_epochs_per_chunk']):
        current_patience += 1
        epoch_losses = []
        epoch_accuracies = []
        batch_times = []
        print(f"Epoch {it + 1:,}/{config['n_epochs_per_chunk']:,}")

        model.train()
        t_epoch_train_start = time()
        for X_ecg, X_cwt, rr, y in tqdm(training_dataloader):
            t_batch_start = time()
            if 'cwt' in config['input_type']:
                X_cwt = X_cwt.to(config['device'])
            else:
                X_cwt = None
            if 'ecg' in config['input_type']:
                X_ecg = X_ecg.to(config['device'])
            else:
                X_ecg = None
            if 'rr' in config['input_type']:
                rr = rr.to(config['device'])
            else:
                rr = None
            y = y.to(torch.int64).to(config['device'])

            y_pred = model(X_ecg, X_cwt, rr)
            loss = criterion(y_pred, y)
            if config['warmup']:
                optimizer.optimizer.zero_grad()
            else:
                optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            y_pred = torch.softmax(y_pred, dim=-1).argmax(dim=-1).detach().cpu()
            bal_acc = balanced_accuracy_score(y.cpu(), y_pred)

            epoch_losses.append(loss.detach().cpu().item())
            epoch_accuracies.append(bal_acc)

            t_batch_end = time()
            batch_times.append(t_batch_end - t_batch_start)
        t_epoch_train_end = time()
        train_loss = (np.mean(epoch_losses), np.std(epoch_losses))
        train_accuracy = (np.mean(epoch_accuracies), np.std(epoch_accuracies))
        train_time = t_epoch_train_end - t_epoch_train_start
        train_batch_time = (np.mean(batch_times), np.std(batch_times))

        with torch.no_grad():
            model.eval()
            epoch_losses = []
            epoch_accuracies = []
            batch_times = []

            t_epoch_eval_start = time()
            for X_ecg, X_cwt, rr, y in tqdm(validation_dataloader):
                t_batch_start = time()
                if 'cwt' in config['input_type']:
                    X_cwt = X_cwt.to(config['device'])
                else:
                    X_cwt = None
                if 'ecg' in config['input_type']:
                    X_ecg = X_ecg.to(config['device'])
                else:
                    X_ecg = None
                if 'rr' in config['input_type']:
                    rr = rr.to(config['device'])
                else:
                    rr = None
                y = y.to(torch.int64).to(config['device'])
                y_pred = model(X_ecg, X_cwt, rr)
                loss = criterion(y_pred, y).detach().cpu().item()

                y_pred = torch.softmax(y_pred, dim=-1).argmax(dim=-1).detach().cpu()
                bal_acc = balanced_accuracy_score(y.cpu(), y_pred)

                epoch_losses.append(loss)
                epoch_accuracies.append(bal_acc)

                t_batch_end = time()
                batch_times.append(t_batch_end - t_batch_start)
            t_epoch_eval_end = time() 
        
        val_loss = (np.mean(epoch_losses), np.std(epoch_losses))
        val_accuracy = (np.mean(epoch_accuracies), np.std(epoch_accuracies))
        val_time = t_epoch_eval_end - t_epoch_eval_start
        val_batch_time = (np.mean(batch_times), np.std(batch_times))

        if val_loss[0] < best_loss:
            current_patience = 0
            if config['model'] == 'cnn':
                torch.save(model, config['checkpoint_path'] / 'cnn.pt')
            else:
                torch.save(model, config['checkpoint_path'] / 'transformer.pt')
            best_loss = val_loss[0]

        print('- Training -')
        print(f"Loss: {train_loss[0]:.4f}±{train_loss[1]:.4f}; Accuracy: {train_accuracy[0]:.2%}±{train_accuracy[1] * 100:.2f}; Time: {train_time:.2f}s - {train_batch_time[0]:.2f}±{train_batch_time[1]:.2f}s/batch")
        print('- Validation -')
        print(f"Loss: {val_loss[0]:.4f}±{val_loss[1]:.4f}; Accuracy: {val_accuracy[0]:.2%}±{val_accuracy[1] * 100:.2f}; Time: {val_time:.2f}s - {val_batch_time[0]:.2f}±{val_batch_time[1]:.2f}s/batch\n")
        log_results(config['checkpoint_path'], chunk_id, it, train_loss[0], train_loss[1], train_accuracy[0], train_accuracy[1],
                    val_loss[0], val_loss[1], val_accuracy[0], val_accuracy[1])
        
    return best_loss, current_patience


def load_datasets(path: Path, input_type: str, verbose=True, cv_fold=None):
    if not cv_fold is None:
        path = path / f"fold_{cv_fold}"
    training_dataset = HDF5Dataset(path / 'train.hdf5')
    validation_dataset = HDF5Dataset(path / 'val.hdf5')

    if verbose:
        print(f"Training using {len(training_dataset.y):,} samples, validating using {len(validation_dataset.y):,} samples\n")
        train_labels, train_counts = np.unique(training_dataset.y, return_counts=True)
        val_labels, val_counts = np.unique(validation_dataset.y, return_counts=True)
        print('Training distribution')
        for l, c in zip(train_labels, train_counts):
            print(f"{l}: {c:,} samples")
        print('\nValidation distribution')
        for l, c in zip(val_labels, val_counts):
            print(f"{l}: {c:,} samples")
        print()

    return training_dataset, validation_dataset


def create_checkpoint(config: Dict[str, Any], iter=0) -> Path:
    checkpoint_path = config['checkpoint_root']
    checkpoint_name = config['checkpoint_name']
    if checkpoint_name is None:
        now = datetime.now()
        now_str = now.strftime('%Y%m%d%H%M')
        checkpoint_dir = checkpoint_path / now_str
    else:
        checkpoint_dir = checkpoint_path / checkpoint_name

    checkpoint_dir = checkpoint_dir / f"iter_{iter}"
    checkpoint_dir.mkdir(exist_ok=True, parents=True)
    create_log_file(checkpoint_dir)
    config['checkpoint_path'] = checkpoint_dir

    with open(checkpoint_dir / 'config.json', 'w') as f:
        json.dump({k: str(v) for k, v in config.items()}, f, indent=4)

    return checkpoint_dir


def initialize_training(config: Dict[str, Any]) -> Tuple[torch.nn.Module, torch.nn.Module, torch.nn.Module]:
    if config['model'] == 'cnn':
        model = CNN2D(n_classes=len(config['allowed_labels']), n_modalities=len(config['input_type'].split('+'))).to(config['device'])
    else:
        model = ECGTransformer(n_input=config['n_input'], input_type=config['input_type'], d_embedding=config['d_embedding'], n_layers=config['n_layers'],
                            n_classes=len(config['allowed_labels']), transformer_dropout=config['dropout'],
                            positional_encoding=config['positional_encoding_scheme'], late_fusion=config['late_fusion']).to(config['device'])
    class_weights = torch.from_numpy(compute_class_weight('balanced', classes=np.unique(training_dataset.y), y=training_dataset.y)).to(torch.float32).to(config['device'])
    criterion = config['criterion'](weight=class_weights if config['weighted_loss'] else None)
    if config['warmup']:
        optimizer = NoamOpt(config['d_embedding'], 1, 800,
            torch.optim.Adam(params=model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9, weight_decay=config['regularization']))
    else:
        optimizer = config['optimizer'](params=model.parameters(), lr=config['lr'], weight_decay=config['regularization'])
    return model, criterion, optimizer


def train(model, criterion, optimizer, training_dataset, validation_dataset, config):
    best_loss = float('inf')
    chunk_id = 0
    current_patience = 0
    while chunk_id < config['n_chunks'] and current_patience < config['patience']:
        chunk_id += 1
        print(f"Processing chunk {chunk_id}; Best loss: {best_loss:.4f}; Current patience: {current_patience} / {config['patience']}")
        training_chunk = load_chunk_of_data(training_dataset, config['chunk_size'], allowed_labels=config['allowed_labels'],
                                            sample_normalization=True, balanced=config['balanced_chunk'])
        validation_chunk = load_chunk_of_data(validation_dataset, config['chunk_size'], allowed_labels=config['allowed_labels'],
                                              sample_normalization=True, balanced=config['balanced_chunk'], validation=True)
        training_chunk_dataset = torch.utils.data.TensorDataset(*training_chunk)
        validation_chunk_dataset = torch.utils.data.TensorDataset(*validation_chunk)

        training_dataloader = torch.utils.data.DataLoader(training_chunk_dataset,
                                                        batch_size=config['batch_size'],
                                                        shuffle=True, num_workers=24)
        validation_dataloader = torch.utils.data.DataLoader(validation_chunk_dataset,
                                                            batch_size=8 * config['batch_size'],
                                                            shuffle=True, num_workers=24)
        best_loss, current_patience = process_chunk(model, criterion, optimizer, training_dataloader, validation_dataloader, config,
                             chunk_id=chunk_id - 1, current_patience=current_patience, best_loss=best_loss)
        
    if current_patience >= config['patience']:
        print("Stopping training due to exceeding patience...")

                
def test(model, testing_dataset, config, chunk_size=1_000, device='cuda'):
    model.eval()

    y_pred = []
    y_true = []
    for chunk_start in range(0, len(testing_dataset), chunk_size):
        print(f"Loading chunk [{chunk_start}:{min(chunk_start + chunk_size, len(testing_dataset))}]...")
        X_ecg, X_cwt, rr, y = testing_dataset[chunk_start:chunk_start + chunk_size]
        y_true.append(y.numpy())

        # Normalize samples
        for i in range(len(y)):
            ecg_sample = X_ecg[i]
            cwt_sample = X_cwt[i]
            nonzero = (ecg_sample[:, -1] < .0).nonzero()
            if len(nonzero) > 0:
                index = nonzero[0]
            else:
                index = len(ecg_sample)
            ecg_sample_min, ecg_sample_max = ecg_sample[:index, :-1].min(), ecg_sample[:index, :-1].max()
            X_ecg[i, :index, :-1] -= ecg_sample_min
            X_ecg[i, :index, :-1] /= ecg_sample_max - ecg_sample_min
            cwt_sample_min, cwt_sample_max = cwt_sample[:index, :-1].min(), cwt_sample[:index, :-1].max()
            X_cwt[i, :index, :-1] -= cwt_sample_min
            X_cwt[i, :index, :-1] /= cwt_sample_max - cwt_sample_min

        with torch.no_grad():
            if 'cwt' in config['input_type']:
                X_cwt = X_cwt.to(device)
            else:
                X_cwt = None
            if 'ecg' in config['input_type']:
                X_ecg = X_ecg.to(device)
            else:
                X_ecg = None
            if 'rr' in config['input_type']:
                rr = rr.to(device)
            else:
                rr = None
            y_pred_chunk = torch.softmax(model(X_ecg, X_cwt, rr), dim=-1).argmax(dim=-1).detach().cpu().numpy()
            y_pred.append(y_pred_chunk)

    y_true = np.concatenate(y_true, axis=0)
    y_pred = np.concatenate(y_pred, axis=0)

    return y_true, y_pred


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Train ECG Transformer')
    for key, value in config.items():
        arg_type = type(value)
        if arg_type is type(Path()):
            arg_type = str
        elif arg_type is type(True):
            arg_type = lambda x: (str(x).lower() == 'true')
        
        parser.add_argument(f'--{key}', type=arg_type, default=value, help=f'Override default value for {key}')

    args = parser.parse_args()

    # Update config with command line arguments
    for key, value in vars(args).items():
        config[key] = value

    if config['model'] == 'cnn':
        config['positional_encoding_scheme'] = None
        config['warmup'] = False
        config['patience'] = 50

    config['data_path'] = Path(config['data_path'])
    config['checkpoint_root'] = Path(config['checkpoint_root'])

    checkpoint_name = config['input_type']
    if 'rr' in config['input_type']:
        checkpoint_name += '_late_fusion' if config['late_fusion'] else '_early_fusion'

    if config['positional_encoding_scheme'] is not None:
        checkpoint_name += f"_{config['positional_encoding_scheme']}"
        
    if len(config['data_path'].name.split('_')) > 3:
        checkpoint_name += f"_{'_'.join(config['data_path'].name.split('_')[3:])}"
    if 'mit' in config['data_path'].name:
        checkpoint_name += '_mit'
    if config['model'] == 'cnn':
        checkpoint_name += '_cnn'

    config['checkpoint_name'] = checkpoint_name

    print(f"Training using the following configuration: {config['input_type']} - {config['positional_encoding_scheme']}")
    
    training_dataset, validation_dataset= load_datasets(config['data_path'], config['input_type'], verbose=True, cv_fold=config['fold_id'])

    for iter in range(config['n_iter']):
        print(f"Starting training iteration {iter + 1} / {config['n_iter']}")
        checkpoint_dir = create_checkpoint(config, iter=iter)
        print(f"Saving checkpoints to {checkpoint_dir}")
        config['checkpoint_path'] = checkpoint_dir
        model, criterion, optimizer = initialize_training(config)
        train(model, criterion, optimizer, training_dataset, validation_dataset, config)

        if config['model'] == 'cnn':
            model = torch.load(config['checkpoint_path'] / 'cnn.pt', weights_only=False)
        else:
            model = torch.load(config['checkpoint_path'] / 'transformer.pt', weights_only=False)
        y_true, y_pred = test(model, validation_dataset, config)
        np.save(config['checkpoint_path'] / 'y_true.npy', y_true)
        np.save(config['checkpoint_path'] / 'y_pred.npy', y_pred)
        print(f"Balanced accuracy: {balanced_accuracy_score(y_true, y_pred):.4f}")

    training_dataset.close()
    validation_dataset.close()



